using SpecialFunctions

function xytors(x,y)
    """
    Purpose : From (x,y) in equilateral triangle to (r,s) coordinates in standard triangle
    """
    L1 = (sqrt(3.0)*y .+ 1.0)/3.0
    L2 = (-3.0*x - sqrt(3.0)*y .+ 2.0)/6.0
    L3 = ( 3.0*x - sqrt(3.0)*y .+ 2.0)/6.0

    r = -L2 + L3 - L1
    s = -L2 - L3 + L1

    r, s
end

function rstoab(r,s)
    Np = length(r); a = zeros(Np,1)

    for n = 1:Np
        if s[n] != 1
            a[n] = 2*(1 .+ r[n]) /(1 .- s[n]) .- 1
        else
            a[n] = -1
        end
    end
    b = s

    a, b
end

function LegendreP(x,N)
    """
    Evaluate orhonormal m'th order Legendre Polynomial at point x
    """
    xp = x
    # Initial values P_0(x) and P_1(x)
    PL = zeros(N+1, length(x))
    PL[1,:]  .= sqrt(1.0/2.0)
    if  N == 0
        P = PL'
        return P
    end
    PL[2,:] = sqrt(3.0/2.0)*xp
    if  N == 1
        P = PL[N+1,:]
        return P
    end

    # Forward recurrence using the symmetry of the recurrence.
    aold = sqrt(1.0/3.0)
    for i=1:N-1
        anew = 2/(2*i+2)*sqrt((i+1)*(i+1)*(i+1)*(i+1)/(2*i+1)/(2*i+3))
        PL[i+2,:] = 1/anew*(-aold*PL[i,:] + xp.*PL[i+1,:])
        aold = anew
    end
    P = PL[N+1,:]
    return P
end

function JacobiP(x,alpha,beta,N)
    """
    Purpose: Evaluate Jacobi Polynomial of type (alpha,beta) > -1
    (alpha+beta <> -1) at points x for order N and returns P[1:length(xp))]
    Note   : They are normalized to be orthonormal.
    """
    xp = x; PL = zeros(N+1, length(x))

    # Initial values P_0(x) and P_1(x)
    gamma0 = 2^(alpha+beta+1)/(alpha+beta+1)*gamma(alpha+1)*gamma(beta+1)/gamma(alpha+beta+1)
    PL[1,:]  .= 1/sqrt(gamma0)
    if  N == 0
        P = PL'
        return P
    end

    gamma1 = (alpha+1)*(beta+1)/(alpha+beta+3)*gamma0
    PL[2,:] = ((alpha+beta+2)*xp/2 .+ (alpha-beta)/2)/sqrt(gamma1)
    if  N == 1
        P = PL[N+1,:]
        return P
    end

    # Repeat value in recurrence
    aold = 2/(2+alpha+beta)*sqrt((alpha+1)*(beta+1)/(alpha+beta+3))

    # Forward recurrence using the symmetry of the recurrence.
    for i=1:N-1
        h1 = 2*i+alpha+beta;
        anew = 2/(h1+2)*sqrt((i+1)*(i+1+alpha+beta)*(i+1+alpha)*(i+1+beta)/(h1+1)/(h1+3))
        bnew = -(alpha^2-beta^2)/h1/(h1+2)
        PL[i+2,:] = 1/anew*( -aold*PL[i,:] .+ (xp .- bnew).*PL[i+1,:])
        aold =anew
    end
    P = PL[N+1,:]
    return P
end

function GradJacobiP(x,alpha,beta,N)
    """
    Purpose: Evaluate the derivative of the Jacobi polynomial of type (alpha,beta)>-1,
    at points r for order N and returns dP[1:length(r))]
    """
    dP = zeros(length(x),1)

    if N == 0
        dP .= 0.0
    else
        dP = sqrt(N*(N+alpha+beta+1))*JacobiP(x,alpha+1,beta+1, N-1)
    end
    dP
end

function Simplex2DP(a,b,i,j)
    """
    Purpose: Evaluate 2D orthonormal polynomial on simplex at (a,b) of order (i,j).
    """
    h1 = LegendreP(a,i)
    h2 = JacobiP(b,2*i+1,0,j)
    P = sqrt(2.0)*h1.*h2.*(1 .- b).^i
    P
end

function GradSimplex2DP(a,b,i,j)
    """
    Purpose: Return the derivatives of the modal basis (i,j) on the 2D simplex at (a,b).
    """
    fa = JacobiP(a, 0, 0, i);     dfa = GradJacobiP(a, 0, 0, i)
    gb = JacobiP(b, 2*i+1,0, j); dgb = GradJacobiP(b, 2*i+1,0, j)

    dmodedr = dfa.*gb
    if i > 0
        dmodedr = dmodedr.*((0.5*(1 .- b)).^(i - 1))
    end

    dmodeds = dfa.*(gb.*(0.5*(1 .+ a)));
    if i > 0
        dmodeds = dmodeds.*((0.5*(1 .- b)).^(i - 1));
    end

    tmp = dgb.*((0.5*(1 .- b)).^i)
    if i > 0
        tmp = tmp - 0.5*i*gb.*((0.5*(1 .- b)).^(i - 1))
    end
    dmodeds = dmodeds+fa.*tmp

    # Normalize
    dmodedr = 2^(i + 0.5)*dmodedr; dmodeds = 2^(i + 0.5)*dmodeds

    dmodedr, dmodeds
end

function Vandermonde1D(m,r)
    """
    Initialize the 1D Vandermonde Matrix, V_{ij} = phi_j(r_i)
    """
    V = zeros(length(r),m+1)

    for j in 1:m+1
        V[:,j] = LegendreP(r, j-1)
    end
    V
end

function Vandermonde2D(N,r,s)
    """
    Purpose : Initialize the 2D Vandermonde Matrix,  V_{ij} = phi_j(r_i, s_i);
    """
    Np = convert(Int, (N+1)*(N+2)/2)
    V = zeros(length(r),Np)

    # Transfer to (a,b) coordinates
    a, b = rstoab(r,s)

    # Build the Vandermonde Matrix
    sk = 1
    for i = 0:N
        for j = 0:N - i
            V[:,sk] = Simplex2DP(a,b,i,j)
            sk = sk + 1
        end
    end
    V
end

function GradVandermonde2D(N,r,s)
    """
    Purpose : Initialize the gradient of the modal basis (i,j) at (r,s) at order N
    """
    Np = convert(Int, (N+1)*(N+2)/2)
    Vr = zeros(length(r), Np); Vs = zeros(length(r), Np)

    # find tensor-product coordinates
    a, b = rstoab(r,s)

    # Initialize Matrix
    index = 1
    for i = 0:N
        for j = 0:N-i
            Vr[:,index], Vs[:,index] = GradSimplex2DP(a,b,i,j)
            index = index + 1
        end
    end
    Vr, Vs
end

function Dmatrices2D(N,r,s,V)
    """
    urpose : Initialize the gradient of the modal basis (i,j) at (r,s) at order N
    """
    Vr, Vs = GradVandermonde2D(N,r,s)
    Dr = Vr/V; Ds = Vs/V
    Dr, Ds
end

function Lift2D(N,Np,Nfaces,Nfp,Fmask,r,s,V)
    """
    Purpose  : Compute surface to volume lift term for DG formulation
    """
    Emat = zeros(Np, Nfaces*Nfp)

    # face 1
    faceR = r[Fmask[:,1]]
    V1D = Vandermonde1D(N, faceR)
    massEdge1 = inv(V1D*V1D')
    Emat[Fmask[:,1],1:Nfp] = massEdge1

    # face 2
    faceR = r[Fmask[:,2]]
    V1D = Vandermonde1D(N, faceR)
    massEdge2 = inv(V1D*V1D')
    Emat[Fmask[:,2],Nfp+1:2*Nfp] = massEdge2

    # face 3
    faceS = s[Fmask[:,3]]
    V1D = Vandermonde1D(N, faceS)
    massEdge3 = inv(V1D*V1D')
    Emat[Fmask[:,3],2*Nfp+1:3*Nfp] = massEdge3

    # inv(mass matrix)*\I_n (L_i,L_j)_{edge_n}
    LIFT = V*(V'*Emat)

    LIFT
end

function EulerDt(N,Q,vmapM,gamma,Fscale)
    """
    Purpose: compute the time step dt for the compressible Euler equations
    """
    rho = Q[:,:,1]; rhou = Q[:,:,2]; rhov = Q[:,:,3]; Ener = Q[:,:,4]
    rho = rho[vmapM]; rhou = rhou[vmapM]; rhov = rhov[vmapM]; Ener = Ener[vmapM]
    u = rhou./rho; v = rhov./rho
    p = (gamma-1.0)*(Ener - rho.*(u.^2 .+ v.^2)/2)
    c = sqrt.(abs.(gamma*p./rho))

    dt = 1.0/maximum(((N+1)^2)*0.5*Fscale[:].*(sqrt.(u[:].^2 .+ v[:].^2 ) .+ c[:]))
    dt
end
